import torch
import torch.nn.functional as F
import matplotlib.pyplot as plt
import random

def cmp(s, dt, t):
    ex = torch.all((dt == t.grad)).item()
    app = torch.allclose(dt, t.grad)
    maxdiff = (dt - t.grad).abs().max().item()
    print(f'{s:15s} | exact: {str(ex):5s} | approximate: {str(app):5s} | maxdiff: {maxdiff}')

def build_dataset(words):
    X, Y = [], []
    for w in words:
            #print(w)
            context = [0] * block_size
            for ch in w + '.':
                idx = stoi[ch]
                X.append(context)
                Y.append(idx)
                #print(''.join(itos[i] for i in context), '--->', itos[idx])
                context = context[1:] + [idx]

    X =   torch.tensor(X)
    Y =   torch.tensor(Y)
    print(X.shape, Y.shape)
    return X, Y

def draw_historgram(Layers):
    plt.figure(figsize=(20,4))
    legends = []
    for i, layer in enumerate(Layers):
        if isinstance(layer, Tanh):
            t = layer.out
            print("layer %d (%10s): mean %+.2f, std %.2f, saturated: %.2f" % (i, layer.__class__, t.mean(), t.std(), (t > 0.97).float().mean() * 100)) 
            hy, hx = torch.histogram(t, density=True)
            plt.plot(hx[:-1].tolist(), hy.tolist())
            legends.append(f'layer {i}')
    plt.legend(legends)
    plt.title("activation histogram")
    plt.show()

def draw_histogram_grad(Layers):
    plt.figure(figsize=(20,4))
    legends = []
    for i, layer in enumerate(Layers):
        if isinstance(layer, Tanh):
            t = layer.out.grad
            print("layer %d (%10s): mean %+.2f, std %.2f" % (i, layer.__class__, t.mean(), t.std()))
            hy, hx = torch.histogram(t, density=True)
            plt.plot(hx[:-1].tolist(), hy.tolist())
            legends.append(f'layer {i}')
    plt.legend(legends)
    plt.title("gradient histogram")
    plt.show()

def draw_2dim_histogram(parameters):
    plt.figure(figsize=(20,4))
    legends = []
    for i, p in enumerate(parameters):
        if p.ndim == 2:
            t = p.grad
            print("weight(%10s): mean %+f, std %.e, grand:data ratio %e" % (tuple(p.shape), t.mean(), t.std(), t.std() / p.std()))
            hy, hx = torch.histogram(t, density=True)
            plt.plot(hx[:-1].tolist(), hy.tolist())
            legends.append(f'{i} {tuple(p.shape)}')
    plt.legend(legends)
    plt.title("gradient histogram")
    plt.show()

def draw_ud_histogram(ud, parameters):
    plt.figure(figsize=(20,4))
    legends = []
    for i, p in enumerate(parameters):
        plt.plot([ud[j][i] for j in range(len(ud))])
        legends.append('param %d' % i)
    plt.plot([0, len(ud)], [-3, -3], 'k') # these ratios shold be ~1e03
    plt.legend(legends)
    plt.title("gradient histogram")
    plt.show()

class Linear:
    def __init__(self, fan_in, fan_out, bias=True, generator=None):
        self.W = torch.randn(fan_in, fan_out, generator=generator) / fan_in**0.5
        self.b = torch.zeros(fan_out) if bias else None
    def __call__(self, x):
        self.out = x @ self.W
        if self.b is not None:
            self.out += self.b
        return self.out
    def parameters(self):
        return [self.W] if self.b is None else [self.W, self.b]
    
class BatchNorm1d:

    def __init__(self, dim, eps=1e-5, momentum=0.1):
        self.eps = eps
        self.momentum = momentum
        self.training = True
        #  parameters (trained with backprop)
        self.gamma = torch.ones(dim)
        self.beta = torch.zeros(dim)
        #  buffers (trained with a running 'momentum')
        self.running_mean = torch.zeros(dim)
        self.running_var = torch.ones(dim)
    def __call__(self, x):
        if self.training:
            xmean = x.mean(0, keepdim=True)
            xvar = x.var(0, keepdim=True)
        else:
            xmean = self.running_mean
            xvar = self.running_var
        xhat = (x - xmean) / torch.sqrt(xvar + self.eps)
        self.out = self.gamma * xhat + self.beta
        if self.training:
            with torch.no_grad():
                self.running_mean = (1 - self.momentum) * self.running_mean + self.momentum * xmean
                self.running_var = (1 - self.momentum) * self.running_var + self.momentum * xvar

        return self.out
    def parameters(self):
        return [self.gamma, self.beta]
    
class Tanh:
    def __init__(self):
        self.out = None
    def __call__(self, x):
        self.out = torch.tanh(x)
        return self.out
    def parameters(self):
        return []
    
if __name__ == '__main__':
    words = open('names.txt', 'r').read().splitlines()
    print(len(words))
    chars = sorted(list(set(''.join(words))))
    chars = ['.'] + chars
    stoi = {s:i for i,s in enumerate(chars)}
    itos = {s:i for i,s in stoi.items()}    
    vocab_size = len(stoi)
    n = len(words)
    n1 = int (0.8*n)
    n2 = int (0.9*n)
    block_size = 3

    X, Y = build_dataset(words[:n1])
    Xval, Yval = build_dataset(words[n1:n2])
    Xtest, Ytest = build_dataset(words[n2:])
    n_embed = 10  # dimension of the embedding
    n_hidden = 200 # number of neurons in hidden layers
    g = torch.Generator().manual_seed(2147483647)
    C = torch.randn((vocab_size, n_embed), generator=g)
    W1  = torch.randn((n_embed * block_size, n_hidden), generator=g) * (5/3) / ((n_embed * block_size) ** 0.5)
    b1  = torch.randn(n_hidden, generator=g) * 0.1
    W2  = torch.randn((n_hidden, vocab_size), generator=g) * 0.1
    b2  = torch.randn(vocab_size, generator=g) * 0.1

    bngain = torch.randn((1, n_hidden), generator=g) * 0.1 + 1.0
    bnbias = torch.randn((1, n_hidden), generator=g) * 0.1

    parameters = [C, W1, b1, W2, b2, bngain, bnbias]
    print(sum(p.nelement() for p in parameters))
    for p in parameters:
        p.requires_grad = True

    batch_size = 32
    n = batch_size
    ix = torch.randint(0, X.shape[0], (batch_size,), generator=g)
    Xb, Yb = X[ix], Y[ix]

    emb = C[Xb]
    embcat = emb.view(emb.shape[0], -1)

    hpredbn = embcat  @ W1 + b1 
    bnmeani = 1/n * hpredbn.sum(0, keepdim=True)
    bndiff = hpredbn - bnmeani
    bndiff2 = bndiff**2
    bnvar = 1/(n-1) * (bndiff2).sum(0, keepdim=True)
    bnvar_inv = (bnvar + 1e-5)**-0.5
    bnraw = bndiff * bnvar_inv
    hpreact = bngain * bnraw + bnbias

    h = torch.tanh(hpreact)

    logits = h @ W2 + b2

    logits_max = logits.max(1, keepdim=True).values
    norm_logits = logits - logits_max
    counts = norm_logits.exp()
    counts_sum = counts.sum(1, keepdim=True)
    counts_sum_inv = counts_sum ** -1
    probs = counts * counts_sum_inv
    logprobs = probs.log()
    loss = -logprobs[range(n), Yb].mean()

    for p in parameters:
        p.grad = None
    
    for t in [logprobs, probs, counts, counts_sum, counts_sum_inv, norm_logits, logits_max, logits, hpreact, h, bnraw, bnvar_inv, bnvar, bndiff2, bndiff, bnmeani, hpredbn, embcat, emb]:
        t.retain_grad()

    loss.backward()
    print(loss)

